import sys
import os
import numpy as np
import matplotlib.pyplot as plt
from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QLabel, QFrame, QScrollArea, QTextEdit, QMessageBox,
    QFileDialog, QSizePolicy
)
from PyQt6.QtCore import Qt, QSize, pyqtSignal
from PyQt6.QtGui import QColor
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qtagg import NavigationToolbar2QT as NavigationToolbar
import seaborn as sns

from aeMode import AEAnalyzer, run_fault_detection
from window_state import manager

class ModernButton(QPushButton):
    def __init__(self, text="", parent=None):
        super().__init__(text, parent)
        self.setFixedHeight(36)
        self.setCursor(Qt.CursorShape.PointingHandCursor)
        self.setStyleSheet("""
            QPushButton {
                background-color: #e6f7ff;
                color: #333333;
                border: 1px solid #91d5ff;
                border-radius: 4px;
                padding: 8px 16px;
                font-size: 14px;
            }
            QPushButton:hover {
                background-color: white;
                border-color: #1890ff;
                color: #1890ff;
            }
            QPushButton:pressed {
                background-color: #40a9ff;
                color: white;
                border-color: #1890ff;
            }
            QPushButton:disabled {
                background-color: #f5f5f5;
                color: #bfbfbf;
                border: 1px solid #d9d9d9;
            }
        """)

class ContentPanel(QFrame):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setStyleSheet("""
            ContentPanel {
                background-color: white;
                border: 1px solid #d9d9d9;
                border-radius: 8px;
                margin: 10px;
            }
        """)
        
        # Create main layout
        self.main_layout = QVBoxLayout(self)
        self.main_layout.setContentsMargins(20, 20, 20, 20)
        self.main_layout.setSpacing(10)
        
        # Header
        self.header = QLabel("自动编码器 (AE) 故障检测工具")
        self.header.setStyleSheet("""
            color: #1890ff;
            font-size: 32px;
            font-weight: bold;
            padding: 20px;
        """)
        self.header.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.main_layout.addWidget(self.header)
        
        # Status Label
        self.status_label = QLabel("请选择数据文件开始分析")
        self.status_label.setStyleSheet("""
            color: #333333;
            font-size: 16px;
            padding: 10px;
        """)
        self.status_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.main_layout.addWidget(self.status_label)
        
        # Progress Text Area
        self.progress_text = QTextEdit()
        self.progress_text.setReadOnly(True)
        self.progress_text.setStyleSheet("""
            QTextEdit {
                background-color: #ffffff;
                border: 1px solid #d9d9d9;
                border-radius: 4px;
                padding: 15px;
                font-family: 'Consolas', 'Monaco', monospace;
                font-size: 14px;
                color: #000000;
                line-height: 1.6;
            }
        """)
        self.progress_text.setFixedHeight(250)  # 增加高度以显示更多内容
        self.progress_text.setAcceptRichText(True)  # 允许富文本
        self.progress_text.hide()  # 初始时隐藏
        self.main_layout.addWidget(self.progress_text)
        
        # 添加训练损失曲线图
        self.figure, self.ax = plt.subplots(figsize=(8, 4))
        self.canvas = FigureCanvas(self.figure)
        self.canvas.setStyleSheet("""
            background-color: white;
            border: 1px solid #d9d9d9;
            border-radius: 4px;
        """)
        self.canvas.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
        self.main_layout.addWidget(self.canvas)
        self.ax.set_title('训练损失曲线', fontsize=16, fontweight='bold')
        self.ax.set_xlabel('训练轮次', fontsize=14)
        self.ax.set_ylabel('损失值', fontsize=14)
        self.ax.grid(True, linestyle='--', alpha=0.7)
        self.canvas.hide()  # 初始隐藏，等到训练开始时显示
        
        # Main content area
        self.content_area = QFrame()
        self.content_area.setStyleSheet("""
            QFrame {
                background-color: white;
                border: 1px solid #d9d9d9;
                border-radius: 8px;
            }
        """)
        self.main_layout.addWidget(self.content_area)
        self.content_area.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
        
        # Content area layout
        self.content_layout = QVBoxLayout(self.content_area)
        self.content_layout.setContentsMargins(20, 20, 20, 20)
        self.content_layout.setSpacing(10)

class AEUI(QMainWindow):
    def __init__(self, main_window=None):
        super().__init__()
        self.main_window = main_window
        
        # 延迟初始化一些变量
        self.analyzer = None
        self.results = None
        self.train_losses = []
        
        # 设置基本窗口属性
        self._setup_window()
        # 创建UI组件
        self._create_ui()
        
    def _setup_window(self):
        """设置窗口基本属性"""
        self.setWindowTitle("AE 故障检测工具")
        self.setStyleSheet("QMainWindow {background-color: #f0f0f0;}")
        
        screen = QApplication.primaryScreen().geometry()
        self.resize(1100, 700)
        self.move(
            (screen.width() - 1100) // 2,
            (screen.height() - 700) // 2
        )
        manager.state_changed.connect(self.on_state_changed)
        self.setWindowState(manager.state)
        
    def _create_ui(self):
        """创建UI组件"""
        # Main widget and layout
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        main_layout = QHBoxLayout(main_widget)
        main_layout.setSpacing(0)
        main_layout.setContentsMargins(0, 0, 0, 0)
        
        # Create sidebar
        sidebar = QWidget()
        sidebar.setFixedWidth(200)
        sidebar.setStyleSheet("""
            QWidget {
                background-color: #f0f0f0;
                border-right: 1px solid #d9d9d9;
            }
        """)
        sidebar_layout = QVBoxLayout(sidebar)
        sidebar_layout.setContentsMargins(10, 20, 10, 20)
        sidebar_layout.setSpacing(10)
        
        # Add buttons to sidebar
        analyze_button = ModernButton("开始分析")
        re2_button = ModernButton("RE2监控图")
        spe_button = ModernButton("SPE监控图")
        
        analyze_button.clicked.connect(self.analyze_data)
        re2_button.clicked.connect(self.show_re2_chart)
        spe_button.clicked.connect(self.show_spe_chart)
        
        sidebar_layout.addWidget(analyze_button)
        sidebar_layout.addWidget(re2_button)
        sidebar_layout.addWidget(spe_button)
        
        # 在按钮区域下方添加伸缩空间
        sidebar_layout.addStretch()
        
        # 添加退出按钮（放在布局底部）
        exit_button = ModernButton("退出主界面")
        exit_button.clicked.connect(self.return_to_main)
        sidebar_layout.addWidget(exit_button)
        
        # Add sidebar to main layout
        main_layout.addWidget(sidebar)
        
        # Add content panel
        self.content = ContentPanel()
        main_layout.addWidget(self.content)

    def update_status(self, message, type='info'):
        color = {
            'info': '#333333',
            'success': '#52c41a',
            'error': '#f5222d',
            'warning': '#faad14'
        }.get(type, '#333333')
        
        self.content.status_label.setStyleSheet(f"""
            color: {color};
            font-size: 16px;
            padding: 10px;
        """)
        self.content.status_label.setText(message)

    def clear_content(self):
        """清除内容区域"""
        for i in reversed(range(self.content.content_layout.count())): 
            widget = self.content.content_layout.itemAt(i).widget()
            if widget is not None:
                widget.deleteLater()

    def load_data(self):
        """加载数据文件"""
        file_path, _ = QFileDialog.getOpenFileName(
            self,
            "选择数据文件",
            "",
            "Excel files (*.xlsx);;CSV files (*.csv)"
        )
        
        if file_path:
            self.file_path = file_path
            self.update_status("数据文件已加载，请点击开始分析", 'success')

    def update_progress(self, message):
        """更新进度信息"""
        # 去除可能的特殊符号
        if message.endswith('...'):
            message = message[:-3]
            
        # 添加文本
        self.content.progress_text.append(message)
        
        # 自动滚动到底部
        self.content.progress_text.verticalScrollBar().setValue(
            self.content.progress_text.verticalScrollBar().maximum()
        )
        QApplication.processEvents()

    def analyze_data(self):
        """分析数据"""
        # 检查是否从主窗口获取数据
        if not hasattr(self, 'file_path'):
            if self.main_window is not None and hasattr(self.main_window, 'file_path') and self.main_window.file_path is not None:
                self.file_path = self.main_window.file_path
            else:
                QMessageBox.warning(self, "警告", "请先在主页加载数据文件")
                return
            
        try:
            # 显示进度文本区域并初始化
            self.content.progress_text.clear()
            self.content.progress_text.show()
            
            # 显示训练损失曲线图
            self.content.canvas.show()
            self.content.ax.clear()
            self.content.ax.set_title('训练损失曲线', fontsize=16, fontweight='bold')
            self.content.ax.set_xlabel('训练轮次', fontsize=14)
            self.content.ax.set_ylabel('损失值', fontsize=14)
            self.content.ax.grid(True, linestyle='--', alpha=0.7)
            self.content.canvas.draw()
            
            # 显示初始化动画
            self.update_progress("🚀 自编码器分析系统初始化...")
            self.update_progress("⚡ 正在准备数据预处理...")
            self.update_progress("-" * 50)  # 分隔线
            self.update_progress("📊 开始数据分析流程")
            self.update_progress("1. 数据加载与标准化")
            self.update_status("正在分析数据...", 'info')
            
            # 获取训练轮数，默认为150
            epochs = getattr(self, 'epochs', 150)
            
            # 创建自定义进度回调函数，用于更新训练损失曲线
            original_update_progress = self.update_progress
            
            def progress_callback_with_plot(message):
                original_update_progress(message)
                
                # 检查消息是否包含训练轮次和损失信息
                if "训练轮次" in message and "损失:" in message:
                    try:
                        # 提取损失值
                        loss_str = message.split("损失:")[1].strip()
                        loss = float(loss_str)
                        
                        # 添加到损失列表
                        self.train_losses.append(loss)
                        
                        # 更新损失曲线
                        epochs = list(range(1, len(self.train_losses) + 1))
                        self.content.ax.clear()
                        self.content.ax.plot(epochs, self.train_losses, 'b-', linewidth=2)
                        self.content.ax.set_title('训练损失曲线', fontsize=16, fontweight='bold')
                        self.content.ax.set_xlabel('训练轮次', fontsize=14)
                        self.content.ax.set_ylabel('损失值', fontsize=14)
                        self.content.ax.grid(True, linestyle='--', alpha=0.7)
                        self.content.canvas.draw()
                    except Exception as e:
                        # 使用更安全的错误处理方式，避免直接print
                        self.update_progress(f"更新损失曲线出错: {str(e)}")
            
            # 运行分析，传递自定义进度回调函数
            self.results = run_fault_detection(self.file_path, progress_callback_with_plot, epochs=epochs)
            
            if self.results:
                self.analyzer = self.results['analyzer']
                self.update_status("分析完成！请选择要查看的图表", 'success')
                
                # 显示分析结果摘要
                self.show_summary()
            else:
                self.update_status("分析失败", 'error')
                
        except Exception as e:
            self.update_status(f"分析出错: {str(e)}", 'error')

    def show_summary(self):
        """显示分析结果摘要"""
        self.clear_content()
        
        summary_text = QTextEdit()
        summary_text.setReadOnly(True)
        summary_text.setStyleSheet("""
            QTextEdit {
                background-color: #f5f5f5;
                border: 1px solid #d9d9d9;
                border-radius: 4px;
                padding: 15px;
                font-family: 'Arial', sans-serif;
                font-size: 14px;
                color: #333333;
                line-height: 1.6;
            }
        """)
        
        # 创建按钮布局
        button_layout = QHBoxLayout()
        
        # 添加查看相关性矩阵按钮
        corr_button = ModernButton("查看相关性矩阵")
        corr_button.clicked.connect(self.show_correlation)
        button_layout.addWidget(corr_button)
        
        # 将按钮布局添加到主布局
        self.content.main_layout.addLayout(button_layout)
        
        # 添加摘要文本
        self.content.main_layout.addWidget(summary_text)
        
        # 填充摘要信息
        if hasattr(self, 'results') and self.results:
            re2_anomalies = self.results.get('re2_anomalies', {})
            spe_anomalies = self.results.get('spe_anomalies', {})
            
            summary_html = f"""
            <h2 style='color: #1890ff;'>分析结果摘要</h2>
            <p><b>RE²检测结果:</b> 检测到 {re2_anomalies.get('count', 0)} 个异常样本 (占比 {re2_anomalies.get('percentage', 0):.2f}%)</p>
            <p><b>SPE检测结果:</b> 检测到 {spe_anomalies.get('count', 0)} 个异常样本 (占比 {spe_anomalies.get('percentage', 0):.2f}%)</p>
            <p><b>模型信息:</b> 自动编码器已成功训练，可用于故障检测</p>
            <p><b>建议:</b> 查看监控图表以获取更详细的分析结果</p>
            """
            
            summary_text.setHtml(summary_html)
        else:
            summary_text.setPlainText("无法获取分析结果摘要")

    def show_correlation(self):
        """显示相关性矩阵"""
        if not self.results:
            QMessageBox.warning(self, "警告", "请先完成数据分析")
            return
            
        self.clear_content()
        try:
            # 添加标题
            title_label = QLabel("变量相关性矩阵")
            title_label.setStyleSheet("""
                font-size: 18px;
                font-weight: bold;
                color: #333;
                margin-bottom: 15px;
            """)
            self.content.content_layout.addWidget(title_label)
            
            # 说明文本
            description = QLabel("下图显示各变量之间的相关性，颜色越深表示相关性越强，红色表示正相关，蓝色表示负相关")
            description.setWordWrap(True)
            description.setStyleSheet("font-size: 14px; color: #555; margin-bottom: 10px;")
            self.content.content_layout.addWidget(description)
            
            # 使用原始数据或相关矩阵绘制相关性热图
            # 检查是否有预先计算的相关矩阵
            if 'correlation_matrix' in self.results and self.results['correlation_matrix'] is not None:
                fig = plt.figure(figsize=(12, 10), dpi=120)
                cmap = sns.diverging_palette(230, 20, as_cmap=True)
                sns.heatmap(self.results['correlation_matrix'],
                          annot=False,
                          cmap=cmap,
                          linewidths=0.5,
                          square=True,
                          vmin=-1, vmax=1,
                          center=0)
                plt.title('变量相关性矩阵', fontsize=18, pad=20, fontweight='bold')
                plt.xlabel('变量', fontsize=14, labelpad=10)
                plt.ylabel('变量', fontsize=14, labelpad=10)
                cbar = plt.gcf().axes[-1]
                cbar.set_ylabel('相关系数', fontsize=12, rotation=270, labelpad=20)
                plt.tight_layout()
            else:
                # 使用训练数据绘制相关性矩阵
                fig = self.analyzer.plot_correlation_matrix(self.results['X_train'])
            
            if fig:
                # 创建高分辨率画布
                canvas = FigureCanvas(fig)
                canvas.setMinimumSize(800, 600)  # 设置画布最小尺寸
                
                # 添加工具栏，支持保存图表等功能
                toolbar = NavigationToolbar(canvas, self)
                
                # 创建图表容器
                chart_container = QWidget()
                chart_layout = QVBoxLayout(chart_container)
                chart_layout.addWidget(toolbar)
                chart_layout.addWidget(canvas)
                
                # 添加到界面
                self.content.content_layout.addWidget(chart_container)
            else:
                self.update_status("无法生成相关性矩阵", 'error')
        except Exception as e:
            self.update_status(f"显示相关性矩阵失败: {str(e)}", 'error')
            import traceback
            traceback.print_exc()

    def return_to_main(self):
        """返回主界面"""
        if self.main_window:
            self.hide()  # 隐藏当前窗口而不是关闭
            # 如果返回主窗口，确保显示主窗口
            self.main_window.show()

    def on_state_changed(self, new_state):
        """响应全局窗口状态改变"""
        if self.windowState() != new_state:
            self.setWindowState(new_state)

    def changeEvent(self, event):
        if event.type() == event.Type.WindowStateChange:
            # 更新全局状态
            manager.state = self.windowState()
        super().changeEvent(event)
        
    # 添加与mainUi.py接口兼容的方法
    def show_analyze(self, params=None):
        """mainUi.py接口兼容方法 - 启动分析
        params: 包含分析参数的字典，例如：
        {
            'epochs': 训练轮数,
            'data_df': 数据DataFrame,
            'file_path': 文件路径
        }
        """
        if params and 'epochs' in params:
            self.epochs = params['epochs']  # 保存轮数设置
        self.analyze_data()
        
    def show_re2_chart(self):
        """显示RE²监控图"""
        self.show_monitoring_chart('re2')
        
    def show_spe_chart(self):
        """显示SPE监控图"""
        self.show_monitoring_chart('spe')

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = AEUI()
    window.show()
    sys.exit(app.exec())